/*
 * Copyright (c) 2015, Freescale Semiconductor, Inc.
 * Copyright 2016-2017 NXP
 * All rights reserved.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include <stdio.h>
#include <string.h>
#include "fsl_sd.h"
#include "fsl_debug_console.h"
#include "ff.h"
#include "diskio.h"
#include "fsl_sd_disk.h"
#include "pin_mux.h"
#include "clock_config.h"
#include "board.h"
#include "sdmmc_config.h"
#include "fsl_common.h"

#include "libtf.h"
/*******************************************************************************
 * Definitions
 ******************************************************************************/
enum {
	BUF = 0,
	STR,
};

/* check if there are some mistakes */
char* errcode[] = {
	[1] = "FR_DISK_ERROR",
	[2 ... 3] = "OTHER_ERROR",
	[4] = "FR_NO_FILE",
	[5] = "FR_NO_PATH",
	[6] = "FR_INVALID_NAME",
	[7 ... 19] = "OTHER_ERROR"
	};

#define FATFS_CHECK(func, decr) \
	error = (func); \
	if(error){ \
		PRINTF(#decr " failed, " "error code %d(%s) ", error, errcode[error]); \
		return -1; \
	}

/* alloc buffer for the read data */
#define READ_BUFFER_ALLOC(ptr, bytesTotal, type)  \
		bytesTotal = f_size(&g_fileObject); \
		ptr = malloc(bytesTotal + type); \
		if(!ptr) { \
			PRINTF("Allocate buffer failed \r\n"); \
			return -1; \
		} \
		if(type == STR) {ptr[bytesTotal] = 0;} /*a str must has a zero ending*/

/* need to free the read buffer after alloc */
#define READ_BUFFER_FREE(ptr) free(ptr);

/* tensor allocate arena */
extern char Image$$RW_TENSOR_ARENA$$Base[];
extern char Image$$RW_TENSOR_ARENA_END$$Base[];
#define TENSOR_ARENA (uint8_t*)Image$$RW_TENSOR_ARENA$$Base
#define TENSOR_ARENA_SIZE (Image$$RW_TENSOR_ARENA_END$$Base - Image$$RW_TENSOR_ARENA$$Base)		
		
#define MAX_LABELS_NAME_LEN 20
	
// from the txt will have a /n, remove it
#define rstrip(str) str[strlen(str)-2] = 0; /* /r/n so -2 */
/*******************************************************************************
 * Prototypes
 ******************************************************************************/
/*!
 * @brief wait card insert function.
 */
static status_t sdcardWaitCardInsert(void);

/*******************************************************************************
 * Variables
 ******************************************************************************/
static FATFS g_fileSystem; /* File system object */
static FIL g_fileObject;   /* File object */
/*******************************************************************************
 * Code
 ******************************************************************************/

/* !
 *	init the fatfs on sdcard 
 */
int fatfs_init();

/*!
 * tflite callback functions 
 */ 
typedef struct py_tf_classify_output_data_callback_data {
    uint32_t item_len;
	float* items;
} tf_classify_output_data_callback_data_t;
static void tf_classify_output_data_callback(void *callback_data,
											 void *model_output,
											 const unsigned int output_height,
											 const unsigned int output_width,
											 const unsigned int output_channels,
											 const bool signed_or_unsigned,
											 const bool is_float)
{
	if((output_height == 1) && (output_width == 1)){
		tf_classify_output_data_callback_data_t* arg = (tf_classify_output_data_callback_data_t*)callback_data;
		int shift = signed_or_unsigned ? 128 : 0;
		arg->item_len = output_channels;
		arg->items = (float*)malloc(sizeof(float) * arg->item_len);
		for (unsigned int i = 0; i < output_channels; i++) {
			if (!is_float) {
				arg->items[i] = ((((uint8_t *) model_output)[i] ^ shift) / 255.0f);
			} else {
				arg->items[i] = (((float *) model_output)[i]);
			}
		}
	}
}

/*!
 * SysTick as timer
 */
volatile uint32_t tick;
void SysTick_Handler(){
	tick++;
}

uint32_t get_ms(){
	return tick;
}

uint64_t get_us(){
	volatile uint64_t tick = (uint64_t)get_ms();
	// if the interval is too short, the us will be 0, and will be cut,
	// because the systick_pass has 24bits, after mul 10^6 will surpass
	// the 32-bits value, so need use 64-bits. Even though we use L to 
	// decalre a value, the compiler will not cast the value
	uint64_t systick_pass = SysTick->LOAD - SysTick->VAL;
	systick_pass *= 1000000L;
	volatile uint64_t us =  systick_pass / SystemCoreClock;
	us += tick * 1000;
	return us;
}
/*!
 * tick callback from libtf 
 */
#define TIME_RESOLUTION get_us
uint64_t time_pass[2];
void tick_callback(){
	static uint32_t enter_num = 0;
	time_pass[enter_num++] = TIME_RESOLUTION();
	enter_num %= 2;
}
/*!
 * @brief Main function
 */
int main(void)
{
    FRESULT error;
    UINT bytesRead;
	uint32_t bytesTotal;
    BOARD_ConfigMPU();
    BOARD_InitBootPins();
    BOARD_BootClockRUN();
    BOARD_InitDebugConsole();
	SysTick_Config(CLOCK_GetCpuClkFreq() / 1000);

    PRINTF("\r\nPlease insert a card into board.\r\n");

	if(fatfs_init()){
		PRINTF("Init fatfs failed. \r\n");
		return -1;
	}

	// first open the model_file names from the model_names.txt
	char* model_name;
	FATFS_CHECK(f_open(&g_fileObject, _T("/model_names.txt"), FA_READ), open model_names.txt);
	READ_BUFFER_ALLOC(model_name, bytesTotal, STR);
	FATFS_CHECK(f_read(&g_fileObject, _T(model_name), bytesTotal, &bytesRead), read mode_names.txt);
	FATFS_CHECK(f_close(&g_fileObject), close model_names.txt);

	// open the model file according the name from model_names.txt
	char* model_buffer;
	FATFS_CHECK(f_open(&g_fileObject, _T(model_name), FA_READ), open model file);
	READ_BUFFER_ALLOC(model_buffer, bytesTotal, BUF);
	FATFS_CHECK(f_read(&g_fileObject, _T(model_buffer), bytesTotal, &bytesRead), read model file);
	FATFS_CHECK(f_close(&g_fileObject), close model file)
	
	// read the image.bin
	char* image_buffer;
	FATFS_CHECK(f_open(&g_fileObject, _T("image.bin"), FA_READ), open image.bin);
	READ_BUFFER_ALLOC(image_buffer, bytesTotal, BUF);
	FATFS_CHECK(f_read(&g_fileObject, _T(image_buffer), bytesTotal, &bytesRead), read image.bin);
	FATFS_CHECK(f_close(&g_fileObject), close image.bin)	
	

	tf_classify_output_data_callback_data_t callback_data;
	uint64_t total_start = TIME_RESOLUTION();
	libtf_invoke((const uint8_t*)model_buffer, TENSOR_ARENA, TENSOR_ARENA_SIZE, tick_callback, image_buffer, tf_classify_output_data_callback, &callback_data);
	uint64_t total_end = TIME_RESOLUTION();
	
	// read the labels.txt and record the top1 output
	uint32_t max_id;
	float max_score = 0.0f;
	uint32_t classify_num = callback_data.item_len;
	char** labels = (char**)malloc(sizeof(char*) * classify_num);
	FATFS_CHECK(f_open(&g_fileObject, _T("labels.txt"), FA_READ), open labels.txt);	
	// merge find the max one and load the labels and 
	for(uint32_t i = 0; i < classify_num; i++){
		float score = callback_data.items[i];
		if(callback_data.items[i] >= max_score){
			max_id = i;
			max_score = score;
		}		
		labels[i] = (char*)malloc(MAX_LABELS_NAME_LEN);
		if(f_gets(labels[i], MAX_LABELS_NAME_LEN, &g_fileObject) == NULL){
			PRINTF("Do not find enough labels, please check labels.txt \r\n");
			break;
		}else{
			rstrip(labels[i]);
		}			
	}
	
	FATFS_CHECK(f_close(&g_fileObject), close labels.txt);	
	
	uint32_t total_costs = total_end - total_start;
	uint32_t model_costs = time_pass[1] - time_pass[0];
	printf("Time costs: %d us(whole process), %d us(only run the model) \r\n", (total_costs), (model_costs));
	printf("Find the best one is %s , score 0.%d\r\n", labels[max_id], (int)(max_score * 100));

	// free the callback_data's items 
	free(callback_data.items);
	callback_data.items = NULL;
	
	// a specific free for the labels
	for(int i=0;i<classify_num;i++){
		free(labels[i]);
	}
	free(labels);


	// free all allocated buffer
	READ_BUFFER_FREE(model_name);
	READ_BUFFER_FREE(model_buffer);
	READ_BUFFER_FREE(image_buffer);
	
	printf("Inference done, waiting restart.... \r\n");

    while (true)
    {
    }
}

static status_t sdcardWaitCardInsert(void)
{
    BOARD_SD_Config(&g_sd, NULL, BOARD_SDMMC_SD_HOST_IRQ_PRIORITY, NULL);

    /* SD host init function */
    if (SD_HostInit(&g_sd) != kStatus_Success)
    {
        PRINTF("\r\nSD host init fail\r\n");
        return kStatus_Fail;
    }

    /* wait card insert */
    if (SD_PollingCardInsert(&g_sd, kSD_Inserted) == kStatus_Success)
    {
        PRINTF("\r\nCard inserted.\r\n");
        /* power off card */
        SD_SetCardPower(&g_sd, false);
        /* power on the card */
        SD_SetCardPower(&g_sd, true);
    }
    else
    {
        PRINTF("\r\nCard detect fail.\r\n");
        return kStatus_Fail;
    }

    return kStatus_Success;
}

int fatfs_init(){
    const TCHAR driverNumberBuffer[3U] = {SDDISK + '0', ':', '/'};
    if (sdcardWaitCardInsert() != kStatus_Success)
    {
        return -1;
    }

    if (f_mount(&g_fileSystem, driverNumberBuffer, 0U))
    {
        PRINTF("Mount volume failed.\r\n");
        return -1;
    }

#if (FF_FS_RPATH >= 2U)
    FRESULT error = f_chdrive((char const *)&driverNumberBuffer[0U]);
    if (error)
    {
        PRINTF("Change drive failed.\r\n");
        return -1;
    }
#endif
	return 0;
}

// in case call the unimplement function
int $Sub$$_sys_open(){
	return 0;
}

